package com.jboltai.capability.model.ollama;

import cn.hutool.core.util.StrUtil;
import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.jboltai.capability.message.AIMessage;
import com.jboltai.capability.message.AIMessageRole;
import com.jboltai.capability.model.AIEventProcessor;
import com.jboltai.config.AIParamKey;
import com.jboltai.config.JBoltAIConfig;
import com.jboltai.event.EventAttrKey;
import com.jboltai.event.EventErrorType;
import com.jboltai.event.EventState;
import com.jboltai.event.ThinkStatus;
import com.jboltai.event.ai.AIEvent;
import com.jboltai.event.ai.aichat.AIChatEvent;
import com.jboltai.resource.ai.AIResource;
import com.jboltai.resource.ai.AIResourceCenter;
import com.jboltai.util.http.HttpUtil;
import com.jboltai.util.http.RequestInstance;
import com.jboltai.util.other.ImgUtil;
import okhttp3.ResponseBody;
import okio.BufferedSource;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.util.List;
import java.util.Objects;

public class OllamaAIImpl implements AIEventProcessor {
    private static final Logger LOGGER = LoggerFactory.getLogger(JBoltAIConfig.LOGGER_NAME);

    public static final OllamaAIImpl INSTANCE = new OllamaAIImpl();

    private static final String STOP_FLAG = "[DONE]";


    @Override
    public void process(AIChatEvent event, AIResource resource) {
        //
        RequestInstance request = buildRequest(event, resource);

        event.setAttr(EventAttrKey.REQUEST_INSTANCE, request);
        LOGGER.debug("开始 调用Ollama大模型 处理AIChat事件：\n\n\nEvent参数: {}, \n\n\nRequest参数: {}", JSON.toJSONString(event), JSON.toJSONString(request));
        processEventByNormal(event, resource, request);
    }

    /**
     * 通过正常请求 处理event
     *
     * @param event
     * @param request
     */
    private static void processEventByNormal(AIChatEvent event, AIResource resource, RequestInstance request) {
        request.addJsonParam(AIParamKey.STREAM, event.getIsStream());
        request.asyncSend(response -> {

            try {
                if (response.code() != 200) {
                    //有异常情况
                    handleFail(event, resource, response.code(), response.body().string(), null);
                    return;
                }

                JSONObject data ;
                if (event.getIsStream()) {
                    //流式请求
                    try (ResponseBody responseBody = response.body(); BufferedSource source = responseBody.source(); BufferedReader reader = new BufferedReader(new InputStreamReader(source.inputStream()))) {

                        String line;
                        while ((line = reader.readLine()) != null) {
                            if (StrUtil.isNotBlank(line)) {
                                data = JSON.parseObject(line);
                                String error = data.getString("error");

                                if (error != null) {
                                    event.fail(EventErrorType.UNKNOWN, error);
                                    break;
                                }



                                handleChatSuccess(event, data, true);
                            }
                        }
                        event.setState(EventState.SUCCESS);

                    }

                } else {
                    //非流式请求
                    String content = response.body().string();
                    data = JSON.parseObject(content);
                    String error = data.getString("error");

                    if (error != null) {
                        event.fail(EventErrorType.UNKNOWN, error);
                        return;
                    }

                    LOGGER.debug("AIChatEvent[{}]请求结束:{}", event.getId(), content);
                    handleChatSuccess(event, data, false);
                    event.setState(EventState.SUCCESS);

                }



            } catch (Exception e) {
                LOGGER.error("AIChatEvent[{}]处理响应结果发生异常:{}", event.getId(), e);
                if (Objects.equals(e.getMessage(), "timeout")) {
                    event.fail(EventErrorType.READ_TIMEOUT, "读取超时，请调整配置中的HttpReadTimeout");
                } else {
                    event.fail(EventErrorType.PROGRAM_EXCEPTION, e.getMessage(), e);
                }
            } finally {
                event.complete();
            }

        }, e -> {
            LOGGER.error("AIChatEvent[{}]请求失败:{}", event.getId(), e);
            event.fail(EventErrorType.NET_ERROR, e);
        });
    }


    /**
     * stream请求时，返回的内容放在"delta" 字段中，非stream请求，返回的内容放在message 字段中
     *
     * @param event
     * @param data
     * @param stream 是否流式
     */
    private static AIMessage handleChatSuccess(AIChatEvent event, JSONObject data, boolean stream) {
        AIMessage message = new AIMessage();
        //获取回复
        JSONObject choice = data.getJSONObject("message");
        if (choice == null) {
            event.fail(EventErrorType.UNKNOWN, "大模型未生成内容");
            return null;
        }
        String content = choice.getString("content");
        content = processThinking(event, message,  content, stream);
        message.setContent(content);
        event.appendResponse(message.getContent());

        message.setRole(AIMessageRole.assistant);
        event.setResultRole(message.getRole());

        //获取使用情况
        Integer promptTokenCount = data.getInteger("prompt_eval_count");
        Integer resTokenCount = data.getInteger("eval_count");
        if (promptTokenCount != null) {
            event.setPromptTokenCount(promptTokenCount);
            event.setTotalTokenCount(promptTokenCount + resTokenCount);
            event.setCompletionTokenCount(resTokenCount);
        }
        if (message.getContent() != null) {
            event.success(message);
        }
        return message;
    }

    /**
     * 处理深度思考内容
     * @param content
     * @param stream
     * @return 返回剔除深度思考后的ai回复内容，null代表目前都是深度思考内容，没有回复内容
     */
    private static String processThinking(AIChatEvent event, AIMessage message, String content, boolean stream) {
        if (stream == false) {
            //非流式
            if (content != null && content.contains("<think>") && content.contains("</think>" )) {
                //提取出思考内容
                int endIndex = content.lastIndexOf("</think>");
                String think = content.substring(content.indexOf("<think>") + 7, endIndex);
                event.thinkOver(think);
                message.setThink(think);
                return content.substring(endIndex + 8); //将深度思考之后的正常回复内容返回
            } else {
                return content;
            }

        }

        //流式
        if (content != null && content.contains("<think>")) {
            String think = content.substring(content.indexOf("<think>") + 7);
            event.thinking(think);
            message.setThink(think);
            return null;
        }

        if (content != null && content.contains("</think>" ) && event.getThinkStatus() == ThinkStatus.RUNNING) {
            String think = content.substring(0, content.indexOf("</think>") );
            event.thinkOver(think);
            message.setThink(think);
            return null;
        }

        if (event.getThinkStatus() == ThinkStatus.RUNNING || event.getThinkStatus() == ThinkStatus.START) {
            event.thinking(content);
            message.setThink(content);
            return null;
        }

        return content;


    }

    /**
     * 失败时的处理
     *
     * @param event
     * @param resource
     * @param code
     * @param responseContent
     * @param exception
     */
    private static void handleFail(AIEvent event, AIResource resource, int code, String responseContent, Throwable exception) {
        JSONObject error = null;
        ;
        LOGGER.error("AIEvent[{}]请求失败:{}_{},异常：{}", event.getId(), code, responseContent, exception);
        if (exception == null) {
            //ai服务异常
            switch (code) {
                case 400: {
                    error = JSON.parseObject(responseContent);
                    event.fail(EventErrorType.REQUEST_PARAM_ERROR, error.getJSONObject("error").getString("message"));
                    break;
                }

                case 401:
                    event.fail(EventErrorType.API_KEY_INVALID, responseContent);
                    break;
                case 429: {
                    error = JSON.parseObject(responseContent);
                    if ("rate_limit".equals(error.getJSONObject("error").getString("code"))) {
                        AIResourceCenter.moveResourceToInvalid(resource);
                        event.fail(EventErrorType.EXCEED_REQUEST_LIMIT, responseContent);
                    } else {
                        event.fail(EventErrorType.ACCOUNT_BALANCE_INSUFFICIENT, responseContent);
                    }
                    break;
                }
                case 500:
                case 503:
                    event.fail(EventErrorType.SERVICE_SERVER_ERROR);
                    break;
                default:
                    //未知错误
                    event.fail(EventErrorType.UNKNOWN, responseContent);
                    break;
            }
        } else {
            //程序异常
            event.fail(EventErrorType.NET_ERROR, exception);
        }
    }


    private RequestInstance buildRequest(AIChatEvent event, AIResource resource) {
        //资源指定了apiurl就用资源的，否则用默认的
        String apiUrl = resource.getApiUrl();
        RequestInstance request = HttpUtil.post(resource.getClient(), apiUrl, event);
        request.setContentTypeToJson();
        request.setAuthorization(resource.getApiKey());
        request.addJsonParam(AIParamKey.MODEL, event.getModelName());

        JSONObject options = new JSONObject();
        if (event.getTemperature() != null) {
            options.put(AIParamKey.TEMPERATURE, event.getTemperature());
        }
        if (event.getTopP() != null) {
            options.put(AIParamKey.TOP_P, event.getTopP());
        }
        if (event.getTopK() != null) {
            options.put(AIParamKey.TOP_K, event.getTopK());
        }
        if (event.getMaxTokens() != null) {
            options.put("num_ctx", event.getMaxTokens());
        }
        if (event.getParams() != null) {
            options.putAll(event.getParams());
        }
        if (!options.isEmpty()) {
            request.addJsonParam("options", options);
        }

        if (event.getMessages() != null) {
            List<AIMessage> list = event.getMessages();

            JSONArray messages = new JSONArray(event.getMessages().size());
            JSONObject item = null;
            //开启思考模式，
            AIMessage promptMsg = list.get(list.size() - 1);
            if (Objects.equals(event.isThinkingEnable(), true)) {
                promptMsg.setContent((promptMsg.getContent() == null ? "" : promptMsg.getContent()) + " /think");
            } else {
                //关闭思考模式， ollama是在提示词后追加 /no_think
                promptMsg.setContent((promptMsg.getContent() == null ? "" : promptMsg.getContent()) + " /no_think");

            }

            for (AIMessage message : list) {
                item = new JSONObject();
                switch (message.getRole()) {
                    case user:
                        item.put("role", "user");
                        break;
                    case assistant:
                        item.put("role", "assistant");
                        break;
                    case system:
                        item.put("role", "system");
                        break;
                    default:
                        continue;
                }
                item.put("content", message.getContent());
                if (message.getImgUrl() != null) {
                    JSONArray imgArray = new JSONArray();
                    //带图片
                    for (String img : message.getImgUrl()) {

                        try {
                            String base64 = ImgUtil.parseToBase64(img);
                            if (StrUtil.isNotBlank(base64)) {
                                imgArray.add(base64.replace(ImgUtil.BASE64_PREFIX, ""));
                            }
                        } catch (Exception e) {
                            LOGGER.error("图片加载失败：{} ", img);
                        }
                    }
                    item.put("images", imgArray);
                }
                messages.add(item);
            }
            request.addJsonParam(AIParamKey.MESSAGES, messages);

        }
        return request;
    }


    @Override
    public void handleCancel(AIChatEvent event, AIResource resource) {
        LOGGER.warn("{}事件取消执行", event.getId());
        RequestInstance requestInstance = event.getAttr(EventAttrKey.REQUEST_INSTANCE);
        requestInstance.cancal();
    }
}
